#include <stdlib.h>
#include <stdio.h>
#include "nbody.h"
#include "video.h"
#include <time.h>
#include <math.h>
#include <string.h>
#include <sys/time.h>
#include "mpi.h"
//include <mpe.h>

double *x, *y; 
double *fx, *fy, *vx, *vy, *m;
int    nbodies;
int    width;
int    height;
int    nsteps;
double g;
double interval;
double min_dist;
int first_row, first_column, noperations, first_body, last_body;
int myrank,nr_proc;
int assignment_strategy;

/* Variables for the MPE logging */
int start_comp, stop_comp;

int    **grid;

/* Reads the parameters and the initial values from the input file */
int init_bodies(char *file) {
     FILE *f; 
     int i;
 
     f = fopen(file,"r");
	
     fscanf(f,"%d %lf %lf %d %lf %d %d\n",&nbodies,&g,&interval,&nsteps,&min_dist,&width,&height);
     x  = malloc(sizeof(double) * nbodies);
     y  = malloc(sizeof(double) * nbodies);
     fx = malloc(sizeof(double) * nbodies);
     fy = malloc(sizeof(double) * nbodies);
     vx = malloc(sizeof(double) * nbodies);
     vy = malloc(sizeof(double) * nbodies);
     m  = malloc(sizeof(double) * nbodies);
     for (i = 0; i < nbodies; i++) {
		fscanf(f,"%lf %lf %lf %lf %lf\n",&x[i],&y[i],&vx[i],&vy[i],&m[i]);
     }
     fclose(f);
     if(myrank == 0) {
     	printf ("NBODY PROBLEM____________________________________________________________________\n");
     	printf ("nbodies: %d\tg var: %0.2lf\tinterval: %0.2lfsec\tnsteps: %d\tmin_dist: %0.2lfp\twidth %d\theight: %d\n", nbodies,g,interval,nsteps,min_dist,width,height);
     }
     return 1;
}

/* Writes the results in the output file */
int finalize_nbodies(char *file) {
     FILE *f;
     int i;
     
     f = fopen(file,"w");
     fprintf(f,"%d %lf %lf %d %lf %d %d\n",nbodies,g,interval, nsteps,min_dist,width,height);
     for (i = 0; i < nbodies; i++) {
                fprintf(f,"%lf %lf %lf %lf %lf\n",x[i],y[i],vx[i],vy[i],m[i]);
     }
     fclose(f);
     return 1;
}

double factorial(double x) {
 double fac = 1;
 int i;
 for (i=2; i<=x; i++) fac *= i;
 return fac;
}

double mypow(double x, double exp) {
 double result = 1;
 int i;

 result = 1;
 for (i=0; i<exp; i++) result *= x;
 return result;
}

double mysin(double alpha) {
      double result = 0;
      int i;

      for (i=0; i < MAX_ACCURACY;i++) {
            result += (double)(mypow(-1,i) / factorial((double)(2*i+1))) * (double)mypow(alpha, 2*i+1);
      }
      return result;
}

double mycos(double alpha) {
      double result = 0;
      int i;

      for (i=0; i < MAX_ACCURACY;i++) {
            result += (double)(mypow(-1,i) / factorial((double)(2*i))) * (double)mypow(alpha, 2*i);
      }
      return result;
}


double myatan(double alpha) {
      double result = 0;
      int i;

      for (i=0; i < MAX_ACCURACY;i++) {
            result += (double)(mypow(-1,i) / (double)(2*i+1)) * (double)mypow(alpha, 2*i+1);
      }
      return result;
}

/* Assigns which operations will execute on each processor to balace the workload across all the processors (Only executed in the load balanced version) */
void assign_load() {
	int remainder, quotient, proc, i, j, total_operations, aux;
	total_operations = 0;
	// Calculates the total amount of operations to compute the forces
	for(i=1; i<nbodies; i++) {
		total_operations += i;
	}
	// Calculates the amount of operations in each processor
	remainder = total_operations % nr_proc;
	quotient = total_operations / nr_proc;
	proc = 0;
	aux = 0;
	// Calculates which operations will execute each processor
	for(i=0; i<nbodies; i++) {
		for(j=i+1; j<nbodies; j++) {
			if(proc < remainder) {
				noperations = quotient + 1;
				if(aux == 0) {
					if(myrank == proc){
						first_row = i;
						first_column = j;
						return;
					}
				}
				aux++;
				if(aux >= noperations) {
					proc++;
					aux = 0;
				}
			}else{
				noperations = quotient;
				if(aux == 0) {
					if(myrank == proc){
						first_row = i;
						first_column = j;
						return;
					}
				}
				aux++;
				if(aux >= noperations) {
					proc++;
					aux = 0;
				}					
			}
		}
	}
}
	

/* Assigns which bodies will be assigned to each processor */	
void assign_bodies() {
	int remainder, quotient, proc, aux;
	remainder = nbodies % nr_proc;
	quotient = nbodies / nr_proc;
	aux = 0;
	for(proc=0; proc<nr_proc; proc++) {
		if(proc < remainder) {
			if(myrank == proc) {
				first_body = aux;
				last_body = aux + quotient;
				return;
			}
			aux += quotient + 1;
		} else {
			if(myrank == proc) {
				first_body = aux;
				last_body = aux + quotient - 1;
				return;
			}
			aux += quotient;
		}
	}
}

/* It calculates the forces exerted on each body by the others as well as the new velocity and position of them (only used in the load balanced version)*/
int calc_nbodies_load() {
     double *ax, *ay, *fxaux, *fyaux, *vxaux, *vyaux, *xaux, *yaux;
     double fax, fay;
     double dist, f, alpha, aux;
     int i,j, count_operations;


     /* Start the calculation of the computation time with MPE */
     //MPE_Log_event(start_comp,0,"start computation (loads balanced)");

     ax = malloc(sizeof(double) * nbodies);
     ay = malloc(sizeof(double) * nbodies);
     fxaux = malloc(sizeof(double) * nbodies);
     fyaux = malloc(sizeof(double) * nbodies);
     vxaux = malloc(sizeof(double) * nbodies);
     vyaux = malloc(sizeof(double) * nbodies);
     xaux = malloc(sizeof(double) * nbodies);
     yaux = malloc(sizeof(double) * nbodies);

     for (i = 0; i < nbodies; i++) {
	fxaux[i] = fyaux[i] = ax[i] = ay[i] = vxaux[i] = vyaux[i] = xaux[i] = yaux[i] = 0;
     }
     count_operations = 0;
     j = first_column;
     // Computes the forces (each processor will calculate only the operations it has assigned)
     for (i = first_row; (i < nbodies) && (count_operations < noperations); i++) {
       for (; (j < nbodies) && (count_operations < noperations); j++) {
          dist = sqrt((x[j] - x[i]) * (x[j] - x[i]) + (y[j] - y[i]) * (y[j] - y[i]));
          if (dist > min_dist ) {
               f = g * m[i] * m[j] / (dist * dist);
               if (f > MAX_F) f = MAX_F;
               aux  = (y[j]-y[i]) / (x[j]-x[i]);
	       alpha = atan(aux);
               fax   = f * mycos(alpha);
	       fay   = f * mysin(alpha);
	       fxaux[i] += fax;
               fyaux[i] += fay;
               fxaux[j] -= fax; 
               fyaux[j] -= fay;
          }
	  count_operations++;
       }
       j = i + 2;
     }

     // Sends the results to the processor with rank = 0 and sums up the values to compute the total forces
     MPI_Reduce (fxaux, fx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
     MPI_Reduce (fyaux, fy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

     // Sends the total forces from the processor with rank = 0 to all the processors
     MPI_Bcast(fx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
     MPI_Bcast(fy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);

     // Calculates the new velocity and position of each body (each processor will calculate only the velocities and positions of the bodies it has assigned)
     for (i = first_body; i <= last_body; i++) {
          ax[i] = fx[i] / m[i];
          ay[i] = fy[i] / m[i];
          vxaux[i] = vx[i] + ax[i] * interval;
          vyaux[i] = vy[i] + ay[i] * interval;
          xaux[i] = x[i] + vxaux[i] * interval;
          yaux[i] = y[i] + vyaux[i] * interval;
          
	  if (xaux[i] <= 0) {
             xaux[i] = 2;
             vxaux[i] = - vxaux[i];
          }
	  if (xaux[i] >= width ) {
             xaux[i] = width - 2;
             vxaux[i] = - vxaux[i];
          }
	  if (yaux[i] <= 0) {
             yaux[i] = 2;
             vyaux[i] = - vyaux[i];
          }
	  if (yaux[i] >= height) {
             yaux[i] = height - 2;
             vyaux[i] = - vyaux[i];
          }
     }

     // Sends the new velocities and positions from each processor to the process with rank = 0 which will gather all the information
     MPI_Reduce (vxaux, vx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
     MPI_Reduce (vyaux, vy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
     MPI_Reduce (xaux, x, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
     MPI_Reduce (yaux, y, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

     // The processor with rank = 0 sends all the new velocities and positions to all the processors
     MPI_Bcast(vx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
     MPI_Bcast(vy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
     MPI_Bcast(x, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
     MPI_Bcast(y, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);

     free(ax);
     free(ay);
     free(fxaux);
     free(fyaux);
     free(vxaux);
     free(vyaux);
     free(xaux);
     free(yaux);

     /* Stop counting time for the computation calculation */
     //MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");

     return 1;
}

/* It calculates the forces exerted on each body by the others as well as the new velocity and position of them (only used in the body balanced version)*/
int calc_nbodies_bodies() {
	double *ax, *ay, *fxaux, *fyaux, *vxaux, *vyaux, *xaux, *yaux;
	double fax, fay;
	double dist, f, alpha, aux;
	int i,j;
	
	/* Start the calculation of the computation time with MPE */
        //MPE_Log_event(start_comp,0,"start computation (bodies balanced)");

	ax = malloc(sizeof(double) * nbodies);
	ay = malloc(sizeof(double) * nbodies);
	fxaux = malloc(sizeof(double) * nbodies);
	fyaux = malloc(sizeof(double) * nbodies);
	vxaux = malloc(sizeof(double) * nbodies);
	vyaux = malloc(sizeof(double) * nbodies);
	xaux = malloc(sizeof(double) * nbodies);
	yaux = malloc(sizeof(double) * nbodies);

	for (i = 0; i < nbodies; i++) {
		fxaux[i] = fyaux[i] = ax[i] = ay[i] = vxaux[i] = vyaux[i] = xaux[i] = yaux[i] = 0;
	}

	// Computes the forces (each processor will calculate only the operations of the bodies it has assigned)
	for (i = first_body; i <= last_body; i++) {
		for (j = i + 1; j < nbodies; j++) {
			dist = sqrt((x[j] - x[i]) * (x[j] - x[i]) + (y[j] - y[i]) * (y[j] - y[i]));
			if (dist > min_dist ) {
				f = g * m[i] * m[j] / (dist * dist);
				if (f > MAX_F) f = MAX_F;
				aux  = (y[j]-y[i]) / (x[j]-x[i]);
				alpha = atan(aux);
				fax   = f * mycos(alpha);
				fay   = f * mysin(alpha);
				fxaux[i] += fax;
				fyaux[i] += fay;
				fxaux[j] -= fax; 
				fyaux[j] -= fay;
			}
		}
	}
	
	// Sends the results to the processor with rank = 0 and sums up the values to compute the total forces
	MPI_Reduce (fxaux, fx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (fyaux, fy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

	// Sends the total forces from the processor with rank = 0 to all the processors
	MPI_Bcast(fx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(fy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);

	// Calculates the new velocity and position of each body (each processor will calculate only the velocities and positions of the bodies it has assigned)
	for (i = first_body; i <= last_body; i++) {
		ax[i] = fx[i] / m[i];
		ay[i] = fy[i] / m[i];
		vxaux[i] = vx[i] + ax[i] * interval;
		vyaux[i] = vy[i] + ay[i] * interval;
		xaux[i] = x[i] + vxaux[i] * interval;
		yaux[i] = y[i] + vyaux[i] * interval;
		
		if (xaux[i] <= 0) {
			xaux[i] = 2;
			vxaux[i] = - vxaux[i];
		}
		if (xaux[i] >= width ) {
			xaux[i] = width - 2;
			vxaux[i] = - vxaux[i];
		}
		if (yaux[i] <= 0) {
			yaux[i] = 2;
			vyaux[i] = - vyaux[i];
		}
		if (yaux[i] >= height) {
			yaux[i] = height - 2;
			vyaux[i] = - vyaux[i];
		}
	}

	// Sends the new velocities and positions from each processor to the process with rank = 0 which will gather all the information
	MPI_Reduce (vxaux, vx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (vyaux, vy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (xaux, x, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (yaux, y, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

	// The processor with rank = 0 sends all the new velocities and positions to all the processors
	MPI_Bcast(vx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(vy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(x, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(y, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);

	free(ax);
	free(ay);
	free(fxaux);
	free(fyaux);
	free(vxaux);
	free(vyaux);
	free(xaux);
	free(yaux);

       /* Stop counting time for the computation calculation */
       //MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");

	return 1;
}
 
int main (int argc, char ** argv) {
     int i, iter;
     int x11 = 0;
     double tstart, tfinish, totalTime;
     iter=0;
     
     if (argc < 5) {
        printf ("Usage:\n");
	printf ("   %s <input_file> <output_file> <X11_0|1> <assignment_strategy>\n", argv[0]);
	printf ("Where:\n");
	printf ("   <X11_0|1> can be 0 or 1. 0 means no X11, 1 means use X11.\n");
	printf ("   <assignment_strategy> can be 0 or 1. 0 means even balance of bodies, 1 means even balance of workload.\n");
        return 0;
     }
     if (!strcmp(argv[3],"1")) x11 = 1;
     if(strcmp(argv[4],"1") == 0) {
	     assignment_strategy = 1;
     } else {
	     if(strcmp(argv[4],"0") == 0) {
		     assignment_strategy = 0;
	     } else {
		     printf ("Usage:\n");
		     printf ("   %s <input_file> <output_file> <X11_0|1> <assignment_strategy>\n", argv[0]);
		     printf ("Where:\n");
		     printf ("   <X11_0|1> can be 0 or 1. 0 means no X11, 1 means use X11.\n");
		     printf ("   <assignment_strategy> can be 0 or 1. 0 means even balance of bodies, 1 means even balance of workload.\n");
		     return 0;
	     }
     }
     /* Initialize MPI */
     MPI_Init(&argc, &argv);
     /* Initialize MPE logging */
     //MPE_Init_log();
     start_comp = MPE_Log_get_event_number();
     stop_comp = MPE_Log_get_event_number();
     //MPE_Describe_state(start_comp,stop_comp,"Computation", "red");
     /* Find out the number of processes */
     MPI_Comm_size(MPI_COMM_WORLD, &nr_proc);
     /* Find out my identity in the default communicator */
     MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
     /* All the processors read the parameters and the initial values from the input file */
     init_bodies(argv[1]);
     /* Start the calculation of the computation time with MPE */
     //MPE_Log_event(start_comp,0,"start computation (bodies balanced)");
     // The process with rank = 0 sets up the window
     if (x11 && myrank == 0) {
	     allocate_grid(width, height);
	     setupWindow(width, height);
     }
     //MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
     // In the load balanced version, each processor calculates which operations will have to compute
     if (assignment_strategy == 1) {
	     assign_load();
     }
     // Each processor calculates which bodies have assigned
     assign_bodies();

     // Each processor executes the body of this loop once on each step of the program
     totalTime = 0;
     for (iter = 0;iter < nsteps;iter++) {
         tstart = MPI_Wtime();
	 // The processors compute the forces, velocities and positions of the bodies
	 if (assignment_strategy == 1) {
		 calc_nbodies_load();
	 } else {
		 calc_nbodies_bodies();
	 }
         tfinish = MPI_Wtime();
         totalTime += tfinish - tstart;
	 // The process with rank = 0 updates the position of the bodies in the window
         if (myrank == 0) {
		 if (x11) {
			 //MPE_Log_event(start_comp,0,"start computation (bodies balanced)");
			 init_grid(width, height);
			 for (i=0; i < nbodies; i++) {
				 grid[(int)x[i]][(int)y[i]]= m[i];
			 }
			 if(iter % REFRESH ==0 ) myDraw(width, height);
			 //MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
		 }
         }
     }

     if (myrank == 0) {
	if (x11) 
		cleanup_grid(width, height);
	finalize_nbodies(argv[2]);
     }
     printf ("Process rank: %d   nbody takes: %fsec    mean: %fsec    accuracy: %d \n", myrank, totalTime, totalTime/nsteps ,MAX_ACCURACY);
     free(x);
     free(y);
     free(fx);
     free(fy);
     free(vx);
     free(vy);
     free(m);
     MPI_Finalize();
     return 1;
}
